"""
辅助函数模块
"""

import os
import torch
import shutil
from config import Config

def save_checkpoint(state, is_best, save_dir):
    """
    保存训练检查点
    """
    cfg = Config()
    filename = os.path.join(save_dir, 'checkpoint.pth.tar')
    torch.save(state, filename)
    if is_best:
        best_filename = os.path.join(save_dir, 'model_best.pth.tar')
        shutil.copyfile(filename, best_filename)

def load_checkpoint(model, optimizer=None, resume_path=None):
    """
    加载检查点
    """
    cfg = Config()
    if resume_path is None:
        resume_path = os.path.join(cfg.SAVE_DIR, 'model_best.pth.tar')
    
    if os.path.isfile(resume_path):
        print(f"=> loading checkpoint '{resume_path}'")
        checkpoint = torch.load(resume_path)
        start_epoch = checkpoint['epoch']
        best_acc = checkpoint['best_acc']
        model.load_state_dict(checkpoint['state_dict'])
        if optimizer is not None:
            optimizer.load_state_dict(checkpoint['optimizer'])
        print(f"=> loaded checkpoint '{resume_path}' (epoch {checkpoint['epoch']}, acc {checkpoint['best_acc']:.2f}%)")
        return start_epoch, best_acc
    else:
        print(f"=> no checkpoint found at '{resume_path}'")
        return 0, 0

def accuracy(output, target, topk=(1,)):
    """
    计算准确率
    """
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res